classednum = 2
import torch.nn as nn

f_dim = 60
class MYLSTM(nn.Module):
    def __init__(self, input_size, hidden_layer_size=10):
        super().__init__()
        self.hidden_layer_size = hidden_layer_size
        self.lstm = nn.LSTM(input_size, hidden_layer_size)
        for p in self.lstm.parameters():
            nn.init.normal_(p,mean=0.0,std=0.001)
        self.linear = nn.Linear(hidden_layer_size, classednum)
        self.sigmoid = nn.Sigmoid()
        self.softmax = nn.Softmax()

    def forward(self, input_x, hidden_cell):
        # print("input_xshape:",input_x.shape)
        input_x = input_x.view(len(input_x), 1, -1)
        # input_x = input_x.view(10, self.bs/10, -1)
        # print("x.shape:",input_x.shape)
        # print(hidden_cell)
        lstm_out, (h_n, h_c) = self.lstm(input_x, hidden_cell)
        # print("lstm_out,h_n,h_c,hidden_cell",lstm_out.shape,h_n.shape,h_c.shape,hidden_cell.shape)
        linear_out = self.linear(lstm_out.view(len(input_x), -1))  # =self.linear(lstm_out[:, -1, :])
        predictions = self.sigmoid(linear_out)
        # predictions = self.softmax(linear_out)
        return predictions, (h_n.data, h_c.data)


class MYRNN(nn.Module):
    def __init__(self, input_size, hidden_size=10):
        super().__init__()
        self.rnn = nn.RNN(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
        )
        for p in self.rnn.parameters():
            nn.init.normal_(p, mean=0.0, std=0.001)
        self.linear = nn.Linear(hidden_size, classednum)
        self.sigmoid = nn.Sigmoid()
        self.hidden_layer_size = hidden_size

    def forward(self, x, hidden_prev):
        input_x = x.view(len(x), 1, -1)
        out, hidden_prev = self.rnn(input_x, hidden_prev)
        # [b, seq, h]
        out = out.view(-1, self.hidden_layer_size)
        out = self.linear(out.view(len(x), -1))  # [seq,h] => [seq,3]
        predictions = self.sigmoid(out)
        return predictions, hidden_prev


class AutoEncoder(nn.Module):
    def __init__(self, inner_coder):
        super(AutoEncoder, self).__init__()
        self.en_conv = nn.Sequential(
            nn.Linear(f_dim, inner_coder),
            # nn.Dropout(0.25),
            # nn.BatchNorm1d(),
            nn.ReLU()
        )
        self.de_conv = nn.Sequential(
            nn.Linear(inner_coder, f_dim),
            nn.ReLU()
        )

    def forward(self, x):
        h = self.en_conv(x)
        decoded = self.de_conv(h)
        # 4 1 8 8
        return h, decoded


class AELinear(nn.Module):
    def __init__(self, initial_dimension, inner_coder):
        super(AELinear, self).__init__()
        self.en_conv = nn.Sequential(
            nn.Linear(initial_dimension, inner_coder),
            # nn.Dropout(0.25),
            # nn.BatchNorm1d(),
            nn.ReLU()
        )
        self.de_conv = nn.Sequential(
            nn.Linear(inner_coder, initial_dimension),
            nn.ReLU()
        )

    def forward(self, x):
        h = self.en_conv(x)
        decoded = self.de_conv(h)
        # 4 1 8 8
        return h, decoded


class AECNN(nn.Module):
    def __init__(self, inner_coder):
        super(AECNN, self).__init__()
        self.prelinear = nn.Linear(f_dim, 64)
        self.en_conv = nn.Sequential(
            # 8x8
            #(8 - 4 + 2*2)/1 + 1
            nn.Conv2d(1, 8, 4, 1, 2),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            # 9x9
            nn.Conv2d(8, 16, 3, 1, 2),
            nn.BatchNorm2d(16),
            nn.ReLU(),
            # 11x11
            nn.Conv2d(16, 10, 4, 2, 2),
            nn.BatchNorm2d(10),
            nn.ReLU(),
            # 6x6  # 11/2 = 5
        )
        self.en_fc = nn.Linear(360, inner_coder)
        self.de_fc = nn.Linear(inner_coder, 360)
        self.de_conv = nn.Sequential(
            # 6x6
            nn.Conv2d(10, 8, 4, 1, 2),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            # 7x7
            nn.Conv2d(8, 4, 4, 1, 2),
            nn.BatchNorm2d(4),
            nn.ReLU(),
            # 8x8
            nn.Conv2d(4, 1, 3, 1, 1),
            nn.Sigmoid(),
            # 8x8
        )
        self.backlinear = nn.Linear(64, f_dim)


    def forward(self, x):
        # print(x.shape)
        x = self.prelinear(x)
        # print(x.shape)
        en = self.en_conv(x.view(x.size(0), 1, 8, 8))
        # print(en.shape)
        en = en.view(x.size(0), -1)
        # print(en.shape)
        code = self.en_fc(en)
        # 4 * 144
        # print(code.shape)
        de = self.de_fc(code)
        # 144
        # print(de.shape)
        decoded = self.de_conv(de.view(x.size(0), 10, 6, 6))
        # print(decoded.shape)
        # 4 1 8 8
        decoded = decoded.reshape(x.size(0), x.size(1))
        # print(decoded.shape)
        decoded = self.backlinear(decoded)
        return code, decoded


class MyCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.prelinear = nn.Linear(f_dim, 64)
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 2, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(2),
            nn.ReLU(),
            nn.Conv2d(2, 4, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(4),
            nn.ReLU(),
            nn.Conv2d(4, 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(),
            nn.MaxPool2d(stride=3, kernel_size=3)
        )
        self.dense = nn.Sequential(
            nn.Linear(32, 32),
            nn.ReLU(),
            nn.Dropout(),
            nn.Linear(32, 10),
            nn.ReLU(),
            nn.Linear(10, classednum)
        )
        self.Sig = nn.Sigmoid()

    def forward(self, input_x):
        input_x = self.prelinear(input_x)
        # print(input_x.shape)
        input_x = input_x.view(input_x.size(0), -1, 8, 8)
        # print(input_x.shape)
        input_x = self.conv1(input_x)
        # print(input_x.shape)
        input_x = input_x.view(1, -1)
        # print(input_x.shape)
        out = self.dense(input_x)
        out.requires_grad_()
        return self.Sig(out)


class MyDNN(nn.Module):
    def __init__(self,input_size):
        super().__init__()
        self.prelinear = nn.Linear(input_size, 64)
        self.d1 = nn.Sequential(
            nn.Linear(64,128),
            nn.ReLU(),
            nn.Linear(128, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512,256),
            nn.ReLU(),
            nn.Linear(256,10),
            nn.Linear(10, classednum)
        )
        self.Sig = nn.Sigmoid()

    def forward(self, input_x):
        input_x = self.prelinear(input_x)
        input_x = self.d1(input_x)
        out = self.Sig(input_x)
        out.requires_grad_()
        return out